// /////////////////////////////////////////////
// OPTIONAL COMPILER OPTIONS             ///////
// /////////////////////////////////////////////

// to use sse and sse2 stuff (requires -O3
// #define SSE2 
// to compute cpu stats
// #define LINUX
// to compile with mpi
// #define PARALLEL


// James N. Simone -- Fermilab

// /////////////////////////////////////////////
// TO USE DOUBLE PRECISION IN ALGORITHMS ///////
// /////////////////////////////////////////////

#if REAL_PRECISION == 4
  #undef USE_DOUBLE_PRECISION
#else
  #ifndef REAL_PRECISION
     #define REAL_PRECISION 8
  #endif
  #define USE_DOUBLE_PRECISION
#endif

// /////////////////////////////////////////////
// LOAD LIBRARIES //////////////////////////////
// /////////////////////////////////////////////

#include "fermiqcd.h"
#include "fermiqcd_JNS_IO.h"
#ifdef SSE2
#include "fermiqcd_fermilab_action.h"
#endif

typedef mdp_real Real;

enum dir   { dirT = 0, dirX, dirY, dirZ, NDIM = 4 };
enum spin  { NSPINS  = 4 };
enum color { NCOLORS = 3 };

enum algorithm { BiCGstab = 0, MinRes };

void print_real_precision ( void )
{
  mpi << "Real precision is set to " <<
      (( sizeof(Real) == sizeof(double) ) ? "double" : "float" )
	      << "\n";
}

// virtual struct
struct source_t
{
  enum source {
    LOCAL,
    WAVEFUNCTION,
    SC_WAVEFUNCTION,
    TIMESLICE,
    EXP_1S,
  };

  typedef enum source source_type;

  virtual void make ( generic_lattice& grid,
		      fermi_field& psi, int spin, int color ) = 0;

  virtual source_type type ( void ) = 0;

  virtual void prompts ( char* file ) = 0;

  virtual ~source_t ( void ) { };
};

struct
{
  int       dim [NDIM];
  Real      kappa; // kappa_s
  Real      zeta;  // kappa_s / kappa_t
  Real      r_t;   // temporal r
  Real      r_s;   // spatial  r
  Real      cSW;   // clover term
  Real      cE;    // electric terms cSW cE
  Real      cB;    // magnetic terms cSW cB
  Real      u0;    // set to 1
  Real      alpha [10]; // alpha [1..9] of improved action
  Real      precision;
  int       max_iterations;
  int       max_restarts;
  int       algorithm;
  source_t* source;
  bool      fix_coul;
  char      fgauge    [FILENAME_MAX];
  char      fgaugefix [FILENAME_MAX];
  char      fquark    [FILENAME_MAX];
}
params;

struct local_source_t : source_t
{
  local_source_t ( void )
    {
      _origin [dirT] = 0;
      _origin [dirX] = 0;
      _origin [dirY] = 0;
      _origin [dirZ] = 0;
    };

  void make ( generic_lattice& grid,
	      fermi_field& psi, int spin, int color )
    {
      site  x ( grid );
      forallsites ( x )
	{
	  psi ( x ) = 0;
	}

      if ( ME == on_which_process ( grid,
				    _origin [dirT],
				    _origin [dirX],
				    _origin [dirY],
				    _origin [dirZ]
				    )
	   )
	{
	  x.set (
		 _origin [dirT],
		 _origin [dirX],
		 _origin [dirY],
		 _origin [dirZ]
		 );

	  // delta function
	  psi ( x, spin, color ) = 1;
	}

      psi.update ( );
    };

  source_type type ( void ) { return LOCAL; };

  void prompts ( char* file )
    {
      string ptr;

      ptr = prompt ( file, "SOURCE-ORIGIN-X", "0" );
      _origin [dirX] = static_cast<int> ( val ( ptr ) );

      ptr = prompt ( file, "SOURCE-ORIGIN-Y", "0" );
      _origin [dirY] = static_cast<int> ( val ( ptr ) );

      ptr = prompt ( file, "SOURCE-ORIGIN-Z", "0" );
      _origin [dirZ] = static_cast<int> ( val ( ptr ) );

      ptr = prompt ( file, "SOURCE-ORIGIN-T", "0" );
      _origin [dirT] = static_cast<int> ( val ( ptr ) );
    };

  int _origin [NDIM];
};

struct exp1S_source_t : source_t
{
  exp1S_source_t ( void ) : _mu ( 1 )
    {
      _origin [dirT] = 0;
      _origin [dirX] = 0;
      _origin [dirY] = 0;
      _origin [dirZ] = 0;
    };

  void make ( generic_lattice& grid, fermi_field& psi, int spin, int color )
    {
      site  s ( grid );
      forallsites ( s )
	{
	  psi ( s ) = 0;
	}

      Real rsq, tmp, norm = 0;
      for ( int x = 0; x < params.dim [dirX]; ++x )
	for ( int y = 0; y < params.dim [dirY]; ++y )
	  for ( int z = 0; z < params.dim [dirZ]; ++z )
	    {

	      if ( ME == on_which_process ( grid,
					    _origin [dirT],
					    x, y, z )
		   )
		{

		  rsq = 0;
		  tmp = x - _origin [dirX];
		  rsq += tmp * tmp;
		  tmp = y - _origin [dirY];
		  rsq += tmp * tmp;
		  tmp = z - _origin [dirZ];
		  rsq += tmp * tmp;
		  rsq = sqrt ( rsq );

		  s.set ( _origin [dirT], x, y, z );

		  tmp = exp ( - _mu * rsq );
		  psi ( s, spin, color ) = Complex ( tmp, 0.0 );
		  norm += tmp * tmp;
		}
	    }
      mpi.add ( norm );

      norm = 1.0 / sqrt ( norm );
      for ( int x = 0; x < params.dim [dirX]; ++x )
	for ( int y = 0; y < params.dim [dirY]; ++y )
	  for ( int z = 0; z < params.dim [dirZ]; ++z )
	    {
	      if ( ME == on_which_process ( grid,
					    _origin [dirT],
					    x, y, z )
		   )
		{
		  s.set ( _origin [dirT], x, y, z );

		  psi ( s, spin, color ) *= Complex ( norm, 0.0 );
		}
	    }

      psi.update ( );
    };

  source_type type ( void ) { return EXP_1S; };

  void prompts ( char* file )
    {
      string ptr;

      ptr = prompt ( file, "SOURCE-MU", "1.00" );
      _mu = val ( ptr );
 
      ptr = prompt ( file, "SOURCE-ORIGIN-X", "0" );
      _origin [dirX] = static_cast<int> ( val ( ptr ) );
 
      ptr = prompt ( file, "SOURCE-ORIGIN-Y", "0" );
      _origin [dirY] = static_cast<int> ( val ( ptr ) );
 
      ptr = prompt ( file, "SOURCE-ORIGIN-Z", "0" );
      _origin [dirZ] = static_cast<int> ( val ( ptr ) );
 
      ptr = prompt ( file, "SOURCE-ORIGIN-T", "0" );
      _origin [dirT] = static_cast<int> ( val ( ptr ) );
     };

  Real _mu;
  int _origin [NDIM];
};

struct wf_source_t : source_t
{
  wf_source_t ( void ) : _origin ( 0 ), _space ( ), _wf ( ),
    _create_cache ( true ), _cache ( )
    {

      int dim [ ] = {
	1,
	params.dim [dirX],
	params.dim [dirY],
	params.dim [dirZ]
      };

      _space.allocate_lattice ( NDIM, dim );
      _wf.allocate_field ( _space, 1 );

      site  s ( _space );
      forallsites ( s )
	{
	  _wf ( s ) = 0;
	}

      _wf.update ( );
    };

  void make ( generic_lattice& grid, fermi_field& psi, int spin, int color )
    {
      site  sf ( grid );
      forallsites ( sf )
	{
	  psi ( sf ) = 0;
	}

      if ( _create_cache )
	{
	  mpi << "caching wavefunction values" << "\n";
	
	  _create_cache = false;

	  // Count wf elements to be cached locally
	  int len = 0;
	  for ( int x = 0; x < params.dim [dirX]; ++x )
	    for ( int y = 0; y < params.dim [dirY]; ++y )
	      for ( int z = 0; z < params.dim [dirZ]; ++z )
		{
		  if ( ME == on_which_process ( grid, _origin, x, y, z ) )
		    {
		      ++len;
		    }
		}

	  if ( len > 0 )
	    {
	      // allocate space
	      _cache.dimension ( 1, len );
	    }

	  // fill the cache
	  int from, to, indx = 0;
	  Complex w;
	  site sw ( _space );
	  for ( int x = 0; x < params.dim [dirX]; ++x )
	    for ( int y = 0; y < params.dim [dirY]; ++y )
	      for ( int z = 0; z < params.dim [dirZ]; ++z )
		{
		  to   = on_which_process ( grid,  _origin, x, y, z );
		  from = on_which_process ( _space,      0, x, y, z );

		  if ( ME == to )
		    {
		      sf.set ( _origin, x, y, z );
		      if ( from == to )
			{
			  sw.set ( 0, x, y, z );

			  _cache ( 0, indx++ ) = _wf ( sw );
			}
		      else
			{
			  mpi.get ( w, from );

			  _cache ( 0, indx++ ) = w;
			}
		    }
		  else if ( ME == from  )
		    {
		      sw.set ( 0, x, y, z );

		      w = _wf ( sw );

		      mpi.put ( w, to );
		    }
		}
	}

      // delete field
      _wf.deallocate_memory ( );

      // do smearing
      int indx = 0;
      for ( int x = 0; x < params.dim [dirX]; ++x )
	for ( int y = 0; y < params.dim [dirY]; ++y )
	  for ( int z = 0; z < params.dim [dirZ]; ++z )
	    {
	      if ( ME == on_which_process ( grid, _origin, x, y, z ) )
		{
		  sf.set ( _origin, x, y, z );
		  psi ( sf, spin, color ) = _cache ( 0, indx++ );
		}
	    }

      psi.update ( );
    };

  source_type type ( void ) { return WAVEFUNCTION; };

  void prompts ( char* file )
    {
      string ptr;

      ptr = prompt ( file, "SOURCE-ORIGIN-T", "0" );
      _origin = static_cast<int> ( val ( ptr ) );

      ptr = prompt ( file, "SOURCE-WAVEFUNCTION", "" );
      _wf_file = ptr;

      unsigned record = 0;
      JNS_read<io_complex<float> > ( _wf, _wf_file.c_str ( ), record,
                                     2 * params.dim [dirX] );
    };

  string                  _wf_file;
  int                     _origin;
  generic_lattice         _space;
  generic_field<Complex>  _wf;
  bool                    _create_cache;
  Matrix                  _cache;
};

void getparams ( int argc, char *argv [ ] )
{
  // /////////////////////////////////////////////
  // read parameters from standard input
  // /////////////////////////////////////////////

  string ptr;

  ptr = prompt ( "<stdin>", "NX", "4" );
  params.dim [dirX] = static_cast<int> ( val ( ptr ) );

  ptr = prompt ( "<stdin>", "NY", "4" );
  params.dim [dirY] = static_cast<int> ( val ( ptr ) );

  ptr = prompt ( "<stdin>", "NZ", "4" );
  params.dim [dirZ] = static_cast<int> ( val ( ptr ) );

  ptr = prompt ( "<stdin>", "NT", "8" );
  params.dim [dirT] = static_cast<int> ( val ( ptr ) );

  ptr = prompt ( "<stdin>", "KAPPA", "0.10" );
  params.kappa = val ( ptr );

  ptr = prompt ( "<stdin>", "R_T", "1.00" );
  params.r_t = val ( ptr );

  ptr = prompt ( "<stdin>", "R_S", "1.00" );
  params.r_s = val ( ptr );

  ptr = prompt ( "<stdin>", "ZETA", "1.00" );
  params.zeta = val ( ptr );

  ptr = prompt ( "<stdin>", "CSW", "1.00" );
  params.cSW = val ( ptr );

  params.cE = 1.0;
  params.cB = 1.0;

  ptr = prompt ( "<stdin>", "ALPHA[1..9]", "0,0,0,0,0,0,0,0,0,0" );

  params.u0 = 1.0;

  // source type
  ptr = prompt ( "<stdin>", "SOURCE", "local" );
  if ( ptr=="local" )
    {
      params.source = new local_source_t ( );
    }
  else if ( ptr=="wavefunction" )
    {
      params.source = new wf_source_t ( );
    }
  else if ( ptr=="exp1S" )
    {
      params.source = new exp1S_source_t ( );
    }
  else
    {
      error("Unimplemented source type");
    };
  params.source -> prompts  ( "<stdin>" );

  ptr = prompt ( "<stdin>", "JIM_GAUGE_INPUT", "cold" );
  strncpy ( params.fgauge, ptr.c_str(), FILENAME_MAX );

  ptr = prompt ( "<stdin>", "GAUGE_FIX_COULOMB", "yes" );
  if ( ptr=="yes" )
    {
      params.fix_coul = true;

      ptr = prompt ( "<stdin>", "JIM_FIXGAUGE_OUTPUT", "none" );
      strncpy ( params.fgaugefix, ptr.c_str(), FILENAME_MAX );
    }
  else
    {
      params.fix_coul = false;

      strcpy ( params.fgaugefix, "none" );
      params.fix_coul = false;
    }

  ptr = prompt ( "<stdin>", "JIM_PROP_OUTPUT", "none" );
  strncpy ( params.fquark, ptr.c_str(), FILENAME_MAX );


  ptr = prompt ( "<stdin>", "SOLVER_ALGORITHM", "BiCGstab" );
  if ( ptr=="MinRes" )
    {
      params.algorithm = MinRes;
    }
  else if ( ptr=="BiCGstab" )
    {
      params.algorithm = BiCGstab;
    }
  else
    {
      mpi << "Unknown algorithm: " << ptr
	  << ", using BiCGstab" << "\n";
  
      params.algorithm = BiCGstab;
    }


  ptr = prompt ( "<stdin>", "PRECISION", "0.5e-12" );
  params.precision = val ( ptr );

  //shutup_loops = TRUE;
  params.max_iterations = 1000;
  params.max_restarts   = 1;
}

struct Pion
{
  int    _nt;
  Matrix _pion;

  Pion ( int nt ) : _nt ( nt ), _pion ( 1, nt )
    {
      clear ( );
    };

  void clear ( void )
    {
      // Zero 2-pt correlator
      for ( int t = 0; t < _nt; ++t )
	{
	  _pion ( 0, t ) = 0.0;
	}
    };

  void show ( void )
    {
      mpi << "Pion propagator" << "\n";
      if ( ME == 0 )
	{
	  for ( int t = 0; t < _nt; ++t )
	    {
	      printf ( "%3i\t%e\n", t, real ( _pion ( 0, t ) ) );
	    }
	}
    };

  void add ( generic_lattice& grid, fermi_field& q )
    {
      Matrix sums ( 1, _nt );
      site x ( grid );

      for ( int t = 0; t < _nt; ++t )
	{
	  sums ( 0, t ) = 0;
	}

      // Pion propagator
      forallsites ( x )
	{
	  sums ( 0, x ( dirT ) ) +=
	    trace ( q ( x ) * hermitian ( q ( x ) ) );
	}

      mpi.add ( sums );
      _pion += sums;
    };
};

int main ( int argc, char *argv [ ] )
{
  int status = 0;

  mpi.open_wormholes ( argc, argv );

  define_base_matrices ( "FERMILAB" );

  print_real_precision ( );

  getparams ( argc, argv );

  // /////////////////////////////////////////////
  // define lattice and fields
  // /////////////////////////////////////////////

  generic_lattice  spacetime ( NDIM, params.dim, default_partitioning0, 
			       torus_topology,0,2,false);

  gauge_field      U   ( spacetime, NCOLORS );
  fermi_field      chi ( spacetime, NCOLORS );
  fermi_field      psi ( spacetime, NCOLORS );

  // pion
  Pion pion ( params.dim [dirT] );

  // /////////////////////////////////////////////
  // select action and inversion method
  // /////////////////////////////////////////////

  coefficients coeff;
  
  coeff["kappa"]  = params.kappa;
  coeff["zeta"]   = params.zeta;
  coeff["r_t"]    = params.r_t;
  coeff["r_s"]    = params.r_s;
  // this is ignored:  coeff["u0"]     = params.u0; 
  coeff["c_{sw}"] = params.cSW;
  coeff["c_E"]    = params.cE;
  coeff["c_B"]    = params.cB;
  coeff["alpha_1"] = params.alpha [1];
  coeff["alpha_2"] = params.alpha [2];
  coeff["alpha_3"] = params.alpha [3];
  coeff["alpha_4"] = params.alpha [4];
  coeff["alpha_5"] = params.alpha [5];
  coeff["alpha_6"] = params.alpha [6];
  coeff["alpha_7"] = params.alpha [7];
  coeff["alpha_8"] = params.alpha [8];
  coeff["alpha_9"] = params.alpha [9];
  
  mpi << coeff << "\n";

  switch ( params.algorithm )
    {
    case MinRes:
      mpi << "using min_residue_inversion" << "\n";
 
      default_fermi_inverter=MinimumResidueInverter<fermi_field,gauge_field>;
      break;

    case BiCGstab:
    default:
      mpi << "using BiCGStab_inversion" << "\n";

      default_fermi_inverter=BiConjugateGradientStabilizedInverter<fermi_field,gauge_field>;
      break;
    }

#ifdef SSE2 
#ifdef FIX_HIGHER_P
  mpi << "using fast fermilab action with sse2" << "\n";

  default_fermi_action     = FermiFermilabActionSSE2::mul_Q;
#else
  mpi << "using fast clover action with sse/sse2" << "\n";
	
  default_fermi_action     = FermiCloverActionSSE2::mul_Q;
#endif

#else

#ifdef FIX_HIGHER_P

  mpi << "Not implemented" << "\n";

#else
  mpi << "using fast clover action without sse" << "\n";
	
  default_fermi_action     = FermiCloverActionFast::mul_Q;
#endif

#endif

  // /////////////////////////////////////////////
  // load gauge configuration (JIM format float)
  // /////////////////////////////////////////////
  if ( strcmp ( "cold", params.fgauge ) == 0 )
    {
      mpi << "Initializing gauge to a cold configuration" << "\n";
      set_cold ( U );
      // JIM_save<float> ( U, params.fgauge );
    }
  else
    {
      unsigned record = 0;
      JNS_read<gl3_links<float> > ( U, params.fgauge,
                                    record, 2 * params.dim [dirX] );
    }

  mpi << "<plaq> = " << average_plaquette ( U ) << "\n";

  // /////////////////////////////////////////////
  // gaugefix (coulomb) and compute propagator
  // /////////////////////////////////////////////
  if ( params.fix_coul )
    {
      int  fix_direction = 0;    // 0 signifies coulomb
      int  fix_sweeps    = 600;
      Real fix_precision = 1e-10; // TODO: canopy 1e-10
      Real fix_overrelax = 1.1;
      bool  fix_z3        = false;
      GaugeFixing::fix ( U, fix_direction, fix_sweeps, fix_precision,
			 fix_overrelax, fix_z3 );

      Real gf_avg_plaq = average_plaquette ( U );
      mpi << "<plaq> = " << gf_avg_plaq << "\n";
	

      if ( strcmp ( "none", params.fgaugefix ) != 0 )
	{
	  mpi << "Saving Coulomb fixed configuration" << "\n";
	  
	  // TODO: must keep gauge meta- separate from quark metadata!!
	  // JIM_save<float> ( U, params.fgaugefix, UserMeta );
	  unsigned record = 0;
	  JNS_read<gl3_links<float> > ( U, params.fgauge,
					record, 2 * params.dim [dirX] );
	}
    }

  if ( params.cSW != 0 )
    compute_em_field ( U );

  inversion_stats stats;

  for ( int spin = 0; spin < NSPINS; ++spin )
    for ( int color = 0; color < NCOLORS; ++color )
      {
	mpi << "# " << ME << ": spin color: " <<
	  spin << " " << color << "\n";

	// ////////////////////////////////////////////
	// create source
	// ////////////////////////////////////////////

	params.source -> make ( spacetime, psi, spin, color );

	// ////////////////////////////////////////////
	// solve
	// ////////////////////////////////////////////

	stats = mul_invQ ( chi, psi, U, coeff, 0, params.precision, params.max_iterations);
	chi.update ( );

	// TODO: check residue in chi or psi? DONE in Fermiqcd2
	if ( stats.relative_precision >= stats.target_relative_precision )
	  {
	    if ( ME == 0 )
	      {
		printf (
		  "No convergence! residue = %.3e > %.3e spin, color %d, %d\n",
		  stats.absolute_precision, params.precision, spin, color );
	      }
	    status = 1;
	    exit ( status );
	  }

	// /////////////////////////////////////////////
	// save propagator (JIM format float)
	// /////////////////////////////////////////////

	// spin-color prpagator file format:
	// 36 byte header ( 9 x unsigned 32-bits )
	// q [src_spin][src_color][t][z][y][x][spin][color]

	if ( strcmp ( "none", params.fquark ) != 0 )
	  {
	    // JIM_save<float> ( chi, params.fquark, UserMeta, spin, color );
	    
            unsigned record = NCOLORS * spin + color;
            JNS_write<Wilson_sc_quark<float> > ( chi, params.fquark,
                                                 record,
                                                 2 * params.dim [dirX]);
	  }

	pion.add ( spacetime, chi );
      }

  pion.show ( );

  // /////////////////////////////////////////////
  // close communications!
  // /////////////////////////////////////////////

  mpi.close_wormholes ( );

  exit ( status );
}
